import numpy as np
import mujoco

from ..robot_state import RobotState


class MJInterface:
    def __init__(self, mj_model: mujoco.MjModel, mj_data: mujoco.MjData):
        super().__init__()
        self.mj_model = mj_model
        self.mj_data = mj_data
        self.timeStep = mj_model.opt.timestep

        self.rpy = np.zeros(3)
        self.baseQuat = np.zeros(4)
        self.f3d = np.zeros((3, 2))
        self.basePos = np.zeros(3)
        self.baseAcc = np.zeros(3)
        self.baseAngVel = np.zeros(3)
        self.baseLinVel = np.zeros(3)
        self.JointName = ['leg_l1_joint', 'leg_l2_joint', 'leg_l3_joint', 'leg_l4_joint', 'leg_l5_joint',
                          'leg_l6_joint', 'leg_r1_joint', 'leg_r2_joint', 'leg_r3_joint', 'leg_r4_joint',
                          'leg_r5_joint', 'leg_r6_joint']
        self.baseName = 'base_link'
        self.orientationSensorName = 'BodyQuat'
        self.velSensorName = 'BodyVel'
        self.gyroSensorName = 'BodyGyro'
        self.accSensorName = 'BodyAcc'

        self.jointNum = len(self.JointName)
        self.jntId_qpos = [-1 for _ in range(self.jointNum)]
        self.jntId_qvel = [-1 for _ in range(self.jointNum)]
        self.jntId_dctl = [-1 for _ in range(self.jointNum)]
        self.motor_pos = np.zeros(self.jointNum)
        self.motor_vel = np.zeros(self.jointNum)
        self.motor_pos_Old = np.zeros(self.jointNum)

        for i in range(self.jointNum):
            tmpId = mujoco.mj_name2id(self.mj_model, mujoco.mjtObj.mjOBJ_JOINT, self.JointName[i])
            self.jntId_qpos[i] = self.mj_model.jnt_qposadr[tmpId]
            self.jntId_qvel[i] = self.mj_model.jnt_dofadr[tmpId]
            motorName = '_'.join(self.JointName[i].split('_')[:-1]) + '_motor'
            tmpId = mujoco.mj_name2id(self.mj_model, mujoco.mjtObj.mjOBJ_ACTUATOR, motorName)
            self.jntId_dctl[i] = tmpId

        self.baseBodyId = mujoco.mj_name2id(self.mj_model, mujoco.mjtObj.mjOBJ_BODY, self.baseName)
        self.orientationSensorId = mujoco.mj_name2id(self.mj_model, mujoco.mjtObj.mjOBJ_SENSOR,
                                                     self.orientationSensorName)
        self.velSensorId = mujoco.mj_name2id(self.mj_model, mujoco.mjtObj.mjOBJ_SENSOR, self.velSensorName)
        self.gyroSensorId = mujoco.mj_name2id(self.mj_model, mujoco.mjtObj.mjOBJ_SENSOR, self.gyroSensorName)
        self.accSensorId = mujoco.mj_name2id(self.mj_model, mujoco.mjtObj.mjOBJ_SENSOR, self.accSensorName)

    def updateSensorValues(self):
        for i in range(self.jointNum):
            self.motor_pos_Old[i] = self.motor_pos[i]
            self.motor_pos[i] = self.mj_data.qpos[self.jntId_qpos[i]]
            self.motor_vel[i] = self.mj_data.qvel[self.jntId_qvel[i]]

        self.baseQuat[:] = self.mj_data.sensor(self.orientationSensorName).data
        tmp = self.baseQuat[0]
        self.baseQuat[0] = self.baseQuat[1]
        self.baseQuat[1] = self.baseQuat[2]
        self.baseQuat[2] = self.baseQuat[3]
        self.baseQuat[3] = tmp

        self.rpy[0] = np.arctan2(2 * (self.baseQuat[3] * self.baseQuat[0] + self.baseQuat[1] * self.baseQuat[2]),
                                 1 - 2 * (self.baseQuat[0] * self.baseQuat[0] + self.baseQuat[1] * self.baseQuat[1]))
        self.rpy[1] = np.arcsin(2 * (self.baseQuat[3] * self.baseQuat[1] - self.baseQuat[0] * self.baseQuat[2]))
        self.rpy[2] = np.arctan2(2 * (self.baseQuat[3] * self.baseQuat[2] + self.baseQuat[0] * self.baseQuat[1]),
                                 1 - 2 * (self.baseQuat[1] * self.baseQuat[1] + self.baseQuat[2] * self.baseQuat[2]))

        posOld = self.basePos.copy()
        self.basePos[:] = self.mj_data.xpos[self.baseBodyId, :]
        self.baseAcc[:] = self.mj_data.sensor(self.accSensorName).data
        self.baseAngVel[:] = self.mj_data.sensor(self.gyroSensorName).data
        self.baseLinVel[:] = (self.basePos - posOld) / self.mj_model.opt.timestep

    def setMotorsTorque(self, tau: np.ndarray):
        for i in range(self.jointNum):
            self.mj_data.ctrl[i] = tau[i]

    def write_state(self, robot_state: RobotState):
        robot_state.motors_pos_cur[:] = self.motor_pos[:]
        robot_state.motors_vel_cur[:] = self.motor_vel[:]
        robot_state.rpy[0: 3] = self.rpy[0: 3]
        robot_state.fL[0: 3] = self.f3d[0: 3, 0]
        robot_state.fR[0: 3] = self.f3d[0: 3, 1]
        robot_state.basePos[0: 3] = self.basePos[0: 3]
        robot_state.baseLinVel[0: 3] = self.baseLinVel[0: 3]
        robot_state.baseAcc[0: 3] = self.baseAcc[0: 3]
        robot_state.baseAngVel[0: 3] = self.baseAngVel[0: 3]
        robot_state.update_state()
